#include "cl_memory_pool.h"

namespace PanNN {
namespace OpenCL {
CLMemeoryPool::CLMemeoryPool(const OpenCLRuntime *openCLRt) : m_openCLRt(openCLRt) {}

CLMemeoryPool::~CLMemeoryPool() {}

cl::Buffer *CLMemeoryPool::CLCreateBuffer(size_t size, cl_mem_flags flags, void *hostData) {
    cl_int errNum = CL_SUCCESS;

    auto bufferID = new (std::nothrow) cl::Buffer(*m_openCLRt->m_context, flags, size, hostData, &errNum);
    if (errNum != CL_SUCCESS) {
        PNN_LOG(ERROR) << "Data Bind Failed, ret = " << errNum;
        return nullptr;
    }

    return bufferID;
}

ErrorCode CLMemeoryPool::CopyDeviceBufferToHost(cl::Buffer *bufferID, void *hostData, size_t size) {
    auto errNum = m_openCLRt->m_defaultCmdQueue->enqueueReadBuffer(*bufferID, CL_TRUE, 0, size, hostData, NULL, NULL);
    return NONE_ERR;
}

ErrorCode CLMemeoryPool::CopyHostToDeviceBuffer(void *hostData, cl::Buffer *bufferID, size_t size) {
    void *deviceData =
        m_openCLRt->m_defaultCmdQueue->enqueueMapBuffer(*bufferID, CL_TRUE, CL_MAP_READ | CL_MAP_WRITE, 0, size);

    if (deviceData != nullptr) {
        ::memcpy(deviceData, hostData, size);
    }

    m_openCLRt->m_defaultCmdQueue->enqueueUnmapMemObject(*bufferID, deviceData);
    return NONE_ERR;
}

// ErrorCode CLMemeoryPool::CopyHostToDeviceSSBO(void *hostData, GLuint ssboBufferID, GLsizeiptr size) {
//     PanNN_ASSERT(m_ssboPool.count(ssboBufferID) > 0);
//     PanNN_ASSERT(m_ssboPool[ssboBufferID].first >= size);

//     glBindBuffer(m_ssboPool[ssboBufferID].second, ssboBufferID);
//     OPENGL_CHECK_ERROR;

//     auto ptr = glMapBufferRange(m_ssboPool[ssboBufferID].second, 0, m_ssboPool[ssboBufferID].first, GL_MAP_READ_BIT);
//     OPENGL_CHECK_ERROR;

//     if (ptr != nullptr) {
//         ::memcpy(ptr, hostData, size);
//     }

//     glUnmapBuffer(m_ssboPool[ssboBufferID].second);
//     OPENGL_CHECK_ERROR;
//     return NONE_ERR;
// }

// GLuint CLMemeoryPool::GLCreateTexture(int w, int h, int c, GLenum textrueFormat, GLenum target) { return textureID; }

cl::Image2D *CLMemeoryPool::CLCreateImage(int w, int h, int c, cl_mem_flags flags, void *hostData) {
    cl_int errNum = CL_SUCCESS;

    flags |= (hostData == nullptr) ? CL_MEM_ALLOC_HOST_PTR : CL_MEM_COPY_HOST_PTR;
    auto imageID = new (std::nothrow)
        cl::Image2D(*m_openCLRt->m_context, flags, cl::ImageFormat(CL_RGBA, CL_FLOAT), w, h, 0, hostData, &errNum);
    if (errNum != CL_SUCCESS) {
        PNN_LOG(ERROR) << "Data Bind Failed, ret = " << errNum;
        return nullptr;
    }

    return imageID;
}

ErrorCode CLMemeoryPool::CopyHostToDeviceImage(void *hostData, cl::Image2D *imageID, const std::vector<size_t> &region,
                                               size_t size) {
    cl::size_type rowPitch;
    cl::size_type slicePitch;
    cl::array<cl::size_type, 3> origin_{0, 0, 0};
    cl::array<cl::size_type, 3> region_{region[0], region[1], region[2]};
    void *deviceData = m_openCLRt->m_defaultCmdQueue->enqueueMapImage(*imageID, CL_TRUE, CL_MAP_READ | CL_MAP_WRITE,
                                                                      origin_, region_, &rowPitch, &slicePitch);

    if (deviceData != nullptr) {
        ::memcpy(deviceData, hostData, size);
    }

    m_openCLRt->m_defaultCmdQueue->enqueueUnmapMemObject(*imageID, deviceData);
    return NONE_ERR;
}

ErrorCode CLMemeoryPool::CopyDeviceImageToHost(void *hostData, cl::Image2D *imageID, const std::vector<size_t> &region,
                                               size_t size) {
    cl::size_type rowPitch;
    cl::size_type slicePitch;
    cl::array<cl::size_type, 3> origin_{0, 0, 0};
    cl::array<cl::size_type, 3> region_{region[0], region[1], region[2]};
    void *deviceData = m_openCLRt->m_defaultCmdQueue->enqueueMapImage(*imageID, CL_TRUE, CL_MAP_READ | CL_MAP_WRITE,
                                                                      origin_, region_, &rowPitch, &slicePitch);

    if (deviceData != nullptr) {
        ::memcpy(hostData, deviceData, size);
    }

    m_openCLRt->m_defaultCmdQueue->enqueueUnmapMemObject(*imageID, deviceData);
    return NONE_ERR;
}

// ErrorCode CLMemeoryPool::CopyDeviceTextureToSSBO(GLuint textureID, GLuint ssboBufferID) {
//     GLuint computeProgram = OpenGLRuntime::CreateComputeProgram(g_glsl_device_to_host);
//     glUseProgram(computeProgram);

//     // bind the src image texture
//     glBindImageTexture(0, textureID, 0, GL_TRUE, 0, GL_READ_ONLY, GL_RGBA32F);

//     // bind the dest output data
//     glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, ssboBufferID);

//     // set uniform values
//     int width = m_texturePool[textureID].first[0];
//     int height = m_texturePool[textureID].first[1];
//     int channel = m_texturePool[textureID].first[2];

//     glUniform1i(2, width);
//     glUniform1i(3, height);
//     glUniform1i(4, channel);

//     int c_4 = UP_DIV(channel, 4);
//     int gLocalSize[3] = {4, 4, 1};
//     glDispatchCompute(UP_DIV(width, gLocalSize[0]), UP_DIV(height, gLocalSize[1]), UP_DIV(c_4, gLocalSize[2]));

//     // memory sync
//     glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
//     glDeleteProgram(computeProgram);
//     return NONE_ERR;
// }

// ErrorCode CLMemeoryPool::CopyDeviceSSBOToTexture(GLuint ssboBufferID, GLuint textureID) {
//     GLuint computeProgram = OpenGLRuntime::CreateComputeProgram(g_glsl_host_to_device);
//     glUseProgram(computeProgram);

//     // bind the src image texture
//     glBindImageTexture(0, textureID, 0, GL_TRUE, 0, GL_WRITE_ONLY, GL_RGBA32F);

//     // bind the dest output data
//     glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, ssboBufferID);

//     // set uniform values
//     int width = m_texturePool[textureID].first[0];
//     int height = m_texturePool[textureID].first[1];
//     int channel = m_texturePool[textureID].first[2];

//     glUniform1i(2, width);
//     glUniform1i(3, height);
//     glUniform1i(4, channel);

//     int c_4 = UP_DIV(channel, 4);
//     int gLocalSize[3] = {4, 4, 1};
//     glDispatchCompute(UP_DIV(width, gLocalSize[0]), UP_DIV(height, gLocalSize[1]), UP_DIV(c_4, gLocalSize[2]));

//     // memory sync
//     glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
//     glDeleteProgram(computeProgram);
//     return NONE_ERR;
// }

// GLuint CLMemeoryPool::CopyHostToDeviceTexture(void *hostData, int width, int height, int channel) {
//     auto ssboBufferID = GLCreateSSBO(sizeof(float) * width * height * channel, hostData);
//     auto textureID = GLCreateTexture(width, height, channel, GL_RGBA32F);
//     CopyDeviceSSBOToTexture(ssboBufferID, textureID);
//     return textureID;
// }

// void *CLMemeoryPool::CopyDeviceTextureToHost(GLuint textureID) {
//     int width = m_texturePool[textureID].first[0];
//     int height = m_texturePool[textureID].first[1];
//     int channel = m_texturePool[textureID].first[2];

//     auto ssboBufferID = GLCreateSSBO(sizeof(float) * width * height * channel);

//     CopyDeviceTextureToSSBO(textureID, ssboBufferID);

//     void *output = malloc(sizeof(float) * width * height * channel);
//     CopyDeviceSSBOToHost(ssboBufferID, output, sizeof(float) * 256);
//     return output;
// }

}  // namespace OpenCL
}  // namespace PanNN
